from Crypto.Cipher import AES
from Crypto.Util.Padding import pad, unpad
import base64


class AESCipher:
    def __init__(self, key, iv):
        self.key = key
        self.iv = iv

    def aes_decrypt(self, encrypted_text):
        """
        AES-CBC 模式解密
        :param encrypted_text: 加密后的文本
        :param key: 密钥(16, 24, 32位)
        :param iv: 初始化向量(16位)
        """
        # 创建 AES 密码对象
        cipher = AES.new(self.key, AES.MODE_CBC, self.iv)

        # 解密
        decrypted = cipher.decrypt(encrypted_text)

        # 去除填充
        # unpadded = unpad(decrypted, AES.block_size)

        return decrypted

    def aes_encrypt(self, text):
        """
        AES-CBC 模式加密
        :param text: 明文
        :param key: 密钥(16, 24, 32位)
        :param iv: 初始化向量(16位)
        """
        # 创建 AES 密码对象
        cipher = AES.new(self.key, AES.MODE_CBC, self.iv)

        # 填充
        padded = pad(text.encode('utf-8'), AES.block_size)

        # 加密
        encrypted = cipher.encrypt(padded)

        return base64.b64encode(encrypted).decode('utf-8')
